import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal

def mahalanobis_distance(p1, p2, cov):
    diff = p1 - p2
    return np.sqrt(diff.T @ np.linalg.inv(cov) @ diff)

# 创建数据点
x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
X, Y = np.meshgrid(x, y)

# 选择两个点
p1 = np.array([0, 0])
p2 = np.array([2, 2])

# 定义协方差矩阵
cov = np.array([[2, 1], [1, 2]])

# 计算马氏距离
pos = np.dstack((X, Y))
rv = multivariate_normal(p1, cov)
Z = rv.pdf(pos)

# 绘制等高线图
plt.figure(figsize=(10, 8))
plt.contourf(X, Y, Z, levels=20, cmap='viridis')
plt.colorbar(label='Probability Density')
plt.plot(p1[0], p1[1], 'ro', markersize=10, label='p1')
plt.plot(p2[0], p2[1], 'bo', markersize=10, label='p2')
plt.plot([p1[0], p2[0]], [p1[1], p2[1]], 'r--', linewidth=2)
plt.title('Mahalanobis Distance Visualization')
plt.xlabel('X')
plt.ylabel('Y')
plt.legend()
plt.show()

print(f"Mahalanobis distance between p1 and p2: {mahalanobis_distance(p1, p2, cov):.2f}")